import sympy
import sympybotics
import scipy.io as sio
import math
import pickle
D1 = 0.2755
D2 = 0.4100
D3 = 0.2073
D4 = 0.0741
D5 = 0.0741
D6 = 0.1600
e2 = 0.0098

_cos = math.cos
_sin = math.sin

aa = math.pi / 6
ca = _cos(aa)
sa = _sin(aa)
c2a = _cos(2*aa)
s2a = _sin(2*aa)
d4b = D3 + (sa/s2a) * D4
d5b = ((sa/s2a) * D4 + (sa/s2a) * D5)
d6b = (sa/s2a) * D5 + D6

rbtdef = sympybotics.RobotDef('jaco', # robot name
                          [('pi/2', 0, D1, 'q'),
                           ('pi', D2, 0, 'q'),
                           ('pi/2', 0, -1*e2, 'q'),
                           ('pi/3', 0, -1*d4b, 'q'),
                           ('pi/3', 0, -1*d5b, 'q'),
                           ('pi', 0, -1*d6b, 'q'),
                           ], # (alpha, a, d, theta)
                            dh_convention='standard' # either 'standard' or 'modified'
                              )
# rbtdef.frictionmodel = {'Coulomb', 'viscous'} # options are None or a combination of 'Coulomb', 'viscous' and 'offset'
rbtdef.dynparms()
rbt = sympybotics.RobotDynCode(rbtdef, verbose=True)
gravity_code = sympybotics.robotcodegen.robot_code_to_func('jl', rbt.g_code, 'gravity_out', 'grav', rbtdef)
f = open('jaco/jaco_gravity_out','w')
f.write(gravity_code)
f.close()

tau_str = sympybotics.robotcodegen.robot_code_to_func('jl', rbt.invdyn_code, 'tau_out', 'tau', rbtdef)
#print(tau_str)
f = open('jaco/jaco_tau_out.cpp','w')
f.write(tau_str)
f.close()

rbt.calc_base_parms()

regressor_str = sympybotics.robotcodegen.robot_code_to_func('jl', rbt.H_code, 'H', 'regressor_func', rbtdef)
#print(regressor_str)
f = open('jaco/jaco_regressor.cpp','w')
f.write(regressor_str)
f.close()

f = open('jaco/jaco_base_para.txt','w')
print(rbt.dyn.n_dynparms)
f.write('the nunber of all parameters is: \n')
f.write(str(rbt.dyn.n_dynparms))
f.write('\n')
f.write('\n')
print(rbt.dyn.dynparms)
f.write('all the parameters is: \n')
f.write(str(rbt.dyn.dynparms))
f.write('\n')
f.write('\n')
print(rbt.dyn.baseparms)
f.write('the base parameters is: \n')
f.write(str(rbt.dyn.baseparms))
f.write('\n')
f.write('\n')
print(rbt.dyn.base_idxs)
f.write('the index of base parameters is: \n')
f.write(str(rbt.dyn.base_idxs))
f.write('\n')
f.write('\n')
print(rbt.dyn.Pb)
f.write('the matrix Pb is: \n')
f.write(str(rbt.dyn.Pb))
f.write('\n')
f.write('\n')
print(rbt.dyn.Pd)
f.write('the matrix Pd is: \n')
f.write(str(rbt.dyn.Pd))
f.write('\n')
f.write('\n')
print(rbt.dyn.Kd)
f.write('the matrix Kd is: \n')
f.write(str(rbt.dyn.Kd))
f.write('\n')
f.write('\n')
f.close
output_file = open('jaco/jaco_regressor_baseparameter.pkl', 'wb')
pickle.dump(rbt, output_file)
output_file1 = open('jaco/jaco_dh.pkl', 'wb')
pickle.dump(rbtdef, output_file1)
